#!/usr/bin/env python

from collections import defaultdict
from operator import itemgetter
import optparse
import os

import numpy

from charmicat import BadCommandLineError, safe_mkdir, Cut

from cmix import fom, parameters, hists

_g_cutvars = 'pcm_pikpi m_pik'.split()                  
_g_byfit_default_args = dict(zip(_g_cutvars,
                                 ({'start': 1.5,
                                   'stop': 2.701,
                                   'step': 0.025,
                                   'namefmt': 'pcm_pikpi_gt_{0:<.3f}',
                                   'cutfmt': 'pcm_pikpi > {0}',
                                   'condvars': {},
                                   },
                                  {'start': 1,
                                   'stop': 10.01,
                                   'step': 0.25,
                                   'namefmt': 'm_pik_not_d0_{0:<.1f}sig',
                                   'cutfmt': '(abs(m_pik - D0_MASS) > {0:.3f}*D0_SIGMA_MC)',
                                   'condvars': parameters.paramdef,
                                   },
                                  )))





def main():
    '''
    Usage: %prog [opts] rsmc.h5 wsmc.h5 var [var2 [...]]
    
    Calculate the FOM for each variable.
    '''

    op = optparse.OptionParser(usage=main.__doc__)
    
    op.add_option("-f", "--force", action="store_true", dest="force", 
                  default=False, help="Force overwrite of output directory")
    op.add_option("-o", "--output", dest='output', default=None,
                  help='Put output in directory DIR', metavar='DIR')

    opts, args = op.parse_args()

    if len(args) < 3:
        msg = "Please supply the right-sign MC PyTables file, "
        msg += "the wrong-sign MC PyTables file, and a cut variable "
        msg = "as arguments to this program."

        raise BadCommandLineError(msg, -1)

    path_rsmc = args.pop(0)
    path_wsmc = args.pop(0)
    cutvars   = args

    for cv in cutvars:
        if cv not in _g_cutvars:       
            msg = "Variable `{0}' not implemented;"
            msg += " choices are {1}".format(cv, _g_cutvars)
            raise BadCommandLineError(msg, -2)

    opts.output = safe_mkdir(opts.output, 
                             default='fom_' + '_'.join(cutvars),
                             force=opts.force)

    best_foms = {}
    foms = defaultdict(list)
    for cv in cutvars:
        cvp = _g_byfit_default_args[cv]
        vals = numpy.arange(cvp['start'], cvp['stop'], cvp['step'])
        for v in vals:
            c = Cut(name=cvp['namefmt'].format(v),
                    pos_str=cvp['cutfmt'].format(v),
                    condvars=cvp['condvars'])
            ifom = fom.fom_fit(c, path_rsmc, path_wsmc, hists['m_kpi'], 
                               plot=os.path.join(opts.output, c.name))
    
            foms[cv].append((v, ifom))
            
        xy = sorted(foms[cv], key=itemgetter(1))

        best_foms[cv] = xy[-1]

    print "Best FOMs:"
    for cv in cutvars:
        print cv, best_foms[cv][0], best_foms[cv][1]
        
    return 0


if __name__ == '__main__':
    try:
        import sys
        sys.exit(main())
    except BadCommandLineError, e:
        import os
        sys.stderr.write(str(e) + "\n")
        msg = "See `{0} -h' for more.\n"
        sys.stderr.write(msg.format(os.path.basename(sys.argv[0])))
        sys.stderr.flush()
        sys.exit(e.retcode)
            
